# -*- coding = utf-8 -*-
# @Author: Vaine0
# @Date: 2020-04-11 16:23:00


import math
import matplotlib.pyplot as plt
from decimal import Decimal
import numpy as np


def main():
    #####输入函数#################
    # 输入基波频率f0
    f0 = Get_f0()  # float
    #print(type(f0), f0)
    # 输入最高谐波次数n
    n = Get_n()  # int
    #print(type(n), n)
    #输入谐波幅度, 相位
    A = []
    b = []
    ip = ""
    for i in range(n):
        print("\nA%d * cos(%d*%s*2πt + b%d)" %
              (i+1, i+1, str(Decimal(str(f0)).to_integral()), i+1))
        A.append(Get_A(i))
        b.append(Get_b(i))
        # print(type(A[0]), A)#float
        # print(type(b[0]), b)#float
        ip = ip + " + %s*cos(%d*%s*2πt+%sπ)" % (str(Decimal(str(A[i])).to_integral()),
                                                i+1, str(Decimal(str(f0)).to_integral()),
                                                str(Decimal(str(b[i])).normalize()))
    print("已输入函数:\n" + ip)
    #####傅里叶变换####################
    input("按下Enter键开始傅里叶变换")
    print("傅里叶变换为:")
    F = FFT(f0, n, A, b)
    print(F)
    #####画图###########################
    SampFreq = 20 * f0 * 2 * math.pi
    # print(SampFreq)
    T = 1/f0
    # print(T)
    t = np.arange(0, 20*T, 1/SampFreq)
    #print(len(t), t)
    f = OriF(f0, n, A, b, t)
    ff = np.arange(0, f0*40/math.log10(f0*10), 1)
    FF = FT(t, f, ff)
    plt.subplots_adjust(left=None, bottom=None, right=None, top=None,
                        wspace=None, hspace=0.5)
    plt.subplot(211)
    plt.plot(t, f)
    plt.xlabel("Time(s)")
    plt.ylabel("f(t)")
    plt.title('Original Wave')
    plt.subplot(212)
    plt.plot(ff, 2*np.abs(FF))
    plt.xlim(-1, f0*40/math.log10(f0*10))
    plt.xlabel("Frequency(Hz)")
    plt.ylabel("F(f)")
    plt.title("Amplitude-Frequency Curve")
    plt.show()

# 输入f0


def Get_f0():
    f0 = input("请输入基波频率:")
    while 1:
        try:
            f0 = float(f0)
            if f0 > 0:
                break
            else:
                int('')
        except:
            print("\n请输入正数!")
            f0 = input("请输入基波频率:")
    return f0
# 输入n


def Get_n():
    n = input("请输入最高谐波次数:")
    while 1:
        try:
            n = int(n)
            if n > 0:
                break
            else:
                int('')
        except:
            print("\n请输入正整数!")
            n = input("请输入最高谐波次数:")
    return n
# 输入Ai


def Get_A(i):
    A = input("请输入%d次谐波的幅度 A%d = " % (i+1, i+1))
    while 1:
        try:
            A = float(A)
            break
        except:
            print("\n请输入实数!")
            A = input("请输入%d次谐波的幅度 A%d = " % (i+1, i+1))
    return A
# 输入bi


def Get_b(i):
    b = input("请输入%d次谐波的初相位 b%d = π * " % (i+1, i+1))
    while 1:
        try:
            b = float(b)
            break
        except:
            print("\n请输入实数!")
            b = input("请输入%d次谐波的初相位 b%d = π * " % (i+1, i+1))
    b = b % 2
    return b
# 原函数


def OriF(f0, n, A, b, t):
    f = 0*t
    for i in range(n):
        f = f + A[i] * np.cos((i+1)*f0*2*math.pi*t + math.pi*b[i])
    #print(len(f), f)
    return f
# 傅里叶变换


def FT(t, f, ff):
    F = np.arange(len(ff), dtype=complex)
    for index in range(len(F)):
        F[index] = complex(np.sum(np.cos(2*np.pi*index*t)*f),
                           -np.sum(np.sin(2*np.pi*index*t)*f))
    return F
# 傅里叶变换


def FFT(f0, n, A, b):
    F = ""
    f0 = str(Decimal(str(f0)).to_integral())
    for i in range(n):
        cosb = str(Decimal(str(math.cos(b[i]))).to_integral())
        sinb = str(Decimal(str(math.sin(b[i]))).to_integral())
        F = F + " + %sπ[(%s + j%s)*δ(w - %s) + (%s - j%s)*δ(w + %s)]" % (str(Decimal(str(A[i])
                                                                                     ).to_integral()), cosb, sinb, str((i+1)*float(f0)), cosb, sinb, str((i+1)*float(f0)))
    return F


if __name__ == "__main__":
    main()
